"""
@Project: pythonPro1
@Name: Canny.py
@Author: linxin_liu
@Date: 2022/10/23 12:02
书124页。
图像f在坐标(x,y)处的梯度为二维列向量：
▽f= ∂f/∂x  f(x+1)-f(x)
    ∂f/∂y  f(y+1)-f(y)
有时 ▽f≈∂f/∂x+∂f/∂y 即 ▽f ≈ f(x+1)-f(x) + f(y+1)-f(y) 但是会损失各向同性。
可以用sobel来计算梯度的幅值。
Canny步骤(527页)：
1. 使用高斯滤波器滤波。
2. 使用 Sobel 滤波器滤波获得在 x 和 y 方向上的输出，在此基础上求出边缘的强度和边缘的角度。
3. 对边缘角度进行量化处理。
4. 根据边缘角度对边缘强度进行非极大值抑制（Non-maximum suppression），使图像边缘变得更细。
5. 使用滞后阈值对图像进行二值化处理，优化图像显示效果。
"""
import cv2
import numpy as np


def Canny(img):
    # 对图像(/灰度图)进行高斯滤波。
    def gaussian_filter(img, K_size=3, sigma=1.4):  # 方差
        if len(img.shape) == 3:
            H, W, C = img.shape
            gray = False
        else:
            img = np.expand_dims(img, axis=-1)
            H, W, C = img.shape
            gray = True

        # 0边界
        pad = K_size // 2
        out = np.zeros([H + pad * 2, W + pad * 2, C], dtype=np.float_)  # 扩大一圈，用来存储结果
        out[pad: pad + H, pad: pad + W] = img.copy().astype(np.float_)  # 原图像复制过去

        # 高斯函数卷积核P525页
        K = np.zeros((K_size, K_size), dtype=np.float_)
        for x in range(-pad, -pad + K_size):
            for y in range(-pad, -pad + K_size):
                K[y + pad, x + pad] = np.exp(- (x ** 2 + y ** 2) / (2 * sigma * sigma))  # e的x幂次方。
        K /= (2 * np.pi * sigma * sigma)
        K /= K.sum()

        tmp = out.copy()

        # filtering
        for y in range(H):
            for x in range(W):
                for c in range(C):
                    out[pad + y, pad + x, c] = np.sum(K * tmp[y: y + K_size, x: x + K_size, c])

        out = np.clip(out, 0, 255)  # 截取
        out = out[pad: pad + H, pad: pad + W]
        out = out.astype(np.uint8)

        if gray:
            out = out[..., 0]

        return out

    # sobel filter
    def sobel_filter(img, K_size=3):
        if len(img.shape) == 3:
            H, W, C = img.shape
        else:
            H, W = img.shape

        # 0边界
        pad = K_size // 2
        out = np.zeros((H + pad * 2, W + pad * 2), dtype=np.float_)
        out[pad: pad + H, pad: pad + W] = img.copy().astype(np.float_)
        tmp = out.copy()

        out_v = out.copy()
        out_h = out.copy()

        # Sobel 垂直
        Kv = [[1., 2., 1.], [0., 0., 0.], [-1., -2., -1.]]
        # Sobel 水平
        Kh = [[1., 0., -1.], [2., 0., -2.], [1., 0., -1.]]

        # filtering
        for y in range(H):
            for x in range(W):
                out_v[pad + y, pad + x] = np.sum(Kv * (tmp[y: y + K_size, x: x + K_size]))
                out_h[pad + y, pad + x] = np.sum(Kh * (tmp[y: y + K_size, x: x + K_size]))

        out_v = np.clip(out_v, 0, 255)
        out_h = np.clip(out_h, 0, 255)

        out_v = out_v[pad: pad + H, pad: pad + W]
        out_v = out_v.astype(np.uint8)
        out_h = out_h[pad: pad + H, pad: pad + W]
        out_h = out_h.astype(np.uint8)

        return out_v, out_h

    # get edge strength and edge angle
    def get_edge_angle(fx, fy):
        # get edge strength 平方
        edge = np.sqrt(np.power(fx.astype(np.float32), 2) + np.power(fy.astype(np.float32), 2))
        edge = np.clip(edge, 0, 255)

        # make sure the denominator is not 0
        fx = np.maximum(fx, 1e-10)
        # fx[np.abs(fx) <= 1e-5] = 1e-5

        # get edge angle
        angle = np.arctan(fy / fx)

        return edge, angle

    # 将角度量化为0°、45°、90°、135°
    def angle_quantization(angle):
        angle = angle / np.pi * 180
        angle[angle < -22.5] = 180 + angle[angle < -22.5]
        _angle = np.zeros_like(angle, dtype=np.uint8)
        _angle[np.where(angle <= 22.5)] = 0
        _angle[np.where((angle > 22.5) & (angle <= 67.5))] = 45
        _angle[np.where((angle > 67.5) & (angle <= 112.5))] = 90
        _angle[np.where((angle > 112.5) & (angle <= 157.5))] = 135

        return _angle

    def non_maximum_suppression(angle, edge):
        H, W = angle.shape
        _edge = edge.copy()

        for y in range(H):
            for x in range(W):
                if angle[y, x] == 0:
                    dx1, dy1, dx2, dy2 = -1, 0, 1, 0
                elif angle[y, x] == 45:
                    dx1, dy1, dx2, dy2 = -1, 1, 1, -1
                elif angle[y, x] == 90:
                    dx1, dy1, dx2, dy2 = 0, -1, 0, 1
                elif angle[y, x] == 135:
                    dx1, dy1, dx2, dy2 = -1, -1, 1, 1
                # 边界处理
                if x == 0:
                    dx1 = max(dx1, 0)
                    dx2 = max(dx2, 0)
                if x == W - 1:
                    dx1 = min(dx1, 0)
                    dx2 = min(dx2, 0)
                if y == 0:
                    dy1 = max(dy1, 0)
                    dy2 = max(dy2, 0)
                if y == H - 1:
                    dy1 = min(dy1, 0)
                    dy2 = min(dy2, 0)
                # 如果不是最大值，则将这个位置像素值置为0
                if max(max(edge[y, x], edge[y + dy1, x + dx1]), edge[y + dy2, x + dx2]) != edge[y, x]:
                    _edge[y, x] = 0

        return _edge

    # 滞后阈值处理二值化图像
    # > HT 的设为255，< LT 的设置0，介于它们两个中间的值，使用8邻域判断法
    def hysterisis(edge, HT=100, LT=30):
        H, W = edge.shape

        # Histeresis threshold
        edge[edge >= HT] = 255
        edge[edge <= LT] = 0

        _edge = np.zeros((H + 2, W + 2), dtype=np.float32)
        _edge[1: H + 1, 1: W + 1] = edge

        # 8 - Nearest neighbor
        nn = np.array(((1., 1., 1.), (1., 0., 1.), (1., 1., 1.)), dtype=np.float32)

        for y in range(1, H + 2):
            for x in range(1, W + 2):
                if _edge[y, x] < LT or _edge[y, x] > HT:
                    continue
                if np.max(_edge[y - 1:y + 2, x - 1:x + 2] * nn) >= HT:
                    _edge[y, x] = 255
                else:
                    _edge[y, x] = 0

        edge = _edge[1:H + 1, 1:W + 1]

        return edge

    # Gray scale 图像融合

    def BGR2GRAY(img):
        b = img[:, :, 0]
        g = img[:, :, 1]
        r = img[:, :, 2]
        # 每个通道不同权值。
        out = 0.2126 * r + 0.7152 * g + 0.0722 * b
        out = out.astype(np.uint8)

        return out

    # grayscale
    gray = BGR2GRAY(img)

    # gaussian filtering
    gaussian = gaussian_filter(gray, K_size=5, sigma=1.4)

    # sobel filtering
    fy, fx = sobel_filter(gaussian, K_size=3)

    # get edge strength, angle
    edge, angle = get_edge_angle(fx, fy)

    # angle quantization
    angle = angle_quantization(angle)

    # non maximum suppression
    edge = non_maximum_suppression(angle, edge)

    # hysterisis threshold
    out = hysterisis(edge, 80, 20)

    return out


if __name__ == '__main__':
    # Read image
    img = cv2.imread("D:/tools/image_operation/bing.png").astype(np.float32)
    # Canny
    edge = Canny(img)
    out = edge.astype(np.uint8)
    # Save result
    cv2.imwrite("D:/tools/image_operation/bing_canny.png", out)
    cv2.imshow("result", out)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
